import random
from typing import Optional, List
import structlog

from app.services.service_discovery import ServiceDiscovery, ServiceHealth

logger = structlog.get_logger()


class LoadBalancer:
    def __init__(self, service_discovery: ServiceDiscovery):
        self.service_discovery = service_discovery
        self.round_robin_counters = {}

    def get_service_instance(self, service_name: str, strategy: str = "round_robin") -> Optional[ServiceHealth]:
        """获取服务实例"""
        service = self.service_discovery.get_service(service_name)

        if not service or not service.is_healthy:
            logger.warning(
                "Service not available",
                service_name=service_name,
                service_exists=service is not None,
                is_healthy=service.is_healthy if service else False
            )
            return None

        return service

    def get_next_instance(self, service_name: str, strategy: str = "round_robin") -> Optional[ServiceHealth]:
        """根据负载均衡策略获取下一个服务实例"""
        healthy_services = self.service_discovery.get_healthy_services()

        if service_name not in healthy_services:
            return None

        service = healthy_services[service_name]

        # 目前只支持轮询策略，可以扩展更多策略
        if strategy == "round_robin":
            return self._round_robin(service_name, service)
        elif strategy == "random":
            return self._random(service_name, service)
        elif strategy == "least_response_time":
            return self._least_response_time(service_name, service)
        else:
            return service

    def _round_robin(self, service_name: str, service: ServiceHealth) -> ServiceHealth:
        """轮询策略"""
        # 由于我们目前只有一个实例，直接返回
        return service

    def _random(self, service_name: str, service: ServiceHealth) -> ServiceHealth:
        """随机策略"""
        return service

    def _least_response_time(self, service_name: str, service: ServiceHealth) -> ServiceHealth:
        """最少响应时间策略"""
        return service